import torch
import random
import numpy as np
from tqdm import tqdm


def set_seed(seed=114514):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def save_checkpoint(state, filename):
    torch.save(state, filename)

def load_checkpoint(filename, state_dict_list):
    checkpoint = torch.load(filename)
    origin_list = []
    for state_dict in state_dict_list:
        origin_list.append(checkpoint[state_dict])
    return origin_list

def check_net_valid(model, optimizer, criterion, loader, epochs, device, accuracy_fn):
    model.train()
    acc_list = []
    loss_list = []
    for epoch in range(1, epochs + 1):
        data, label = next(iter(loader))
        data, label = data.to(device), label.to(device)

        # forward
        output = model(data)
        loss = criterion(output, label)
        accuracy = accuracy_fn(output, label)

        # backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        acc_list.append(accuracy.item())
        loss_list.append(loss.item())

    mean_loss = sum(loss_list) / len(loss_list)
    mean_acc = sum(acc_list) / len(acc_list)

    print(f"Train/Loss: {mean_loss}")
    print(f"Train/Accuracy: {mean_acc}")



def train(model, optimizer, scheduler, criterion, train_loader, valid_loader, epochs, device,
          accuracy_fn, max_patience, writer=None):
    model.train()

    patience = 0
    max_accuracy = 0

    scalar = torch.amp.GradScaler('cuda')

    for epoch in range(1, epochs + 1):
        train_tqdm_loader = tqdm(train_loader,desc=f"Train: Epoch [{epoch}/{epochs}]")

        acc_list = []
        loss_list = []


        for data, label in train_tqdm_loader:
            data, label = data.to(device), label.to(device)

            # forward
            with torch.amp.autocast('cuda'):
                output = model(data)
                loss = criterion(output, label)
                accuracy = accuracy_fn(output, label)

            # backward
            optimizer.zero_grad()
            scalar.scale(loss).backward()
            # loss.backward()
            # optimizer.step()
            scalar.step(optimizer)
            scalar.update()


            acc_list.append(accuracy.item())
            loss_list.append(loss.item())

            train_tqdm_loader.set_postfix(loss=loss.item(), accuracy=accuracy.item())

        mean_loss = sum(loss_list) / len(loss_list)
        mean_acc = sum(acc_list) / len(acc_list)
        scheduler.step(mean_loss)

        tqdm.write(f"Epoch [{epoch}] Train/Loss: {mean_loss:.5}")
        tqdm.write(f"Epoch [{epoch}] Train/Accuracy: {mean_acc:.5}")
        if writer is not None:
            writer.add_scalar("Train/Loss", mean_loss, epoch)
            writer.add_scalar("Train/Accuracy", mean_acc, epoch)

        valid_loss, valid_acc = test(model, criterion, valid_loader, device, accuracy_fn, mode="Valid")
        if valid_acc > max_accuracy:
            max_accuracy = valid_acc
            patience = 0
        else:
            patience += 1
            if patience == max_patience:
                break


def test(model, criterion, test_loader, device, accuracy_fn, mode="Test"):
    model.eval()
    with torch.no_grad():
        tqdm_loader = tqdm(test_loader, desc=f"{mode}: ")
        acc_list = []
        loss_list = []

        for data, label in tqdm_loader:
            data, label = data.to(device), label.to(device)

            output = model(data)
            loss = criterion(output, label)
            accuracy = accuracy_fn(output, label)

            acc_list.append(accuracy.item())
            loss_list.append(loss.item())

            tqdm_loader.set_postfix(loss=loss.item(), accuracy=accuracy.item(), )

        mean_loss = sum(loss_list) / len(loss_list)
        mean_acc = sum(acc_list) / len(acc_list)
        tqdm.write(f"{mode}/Loss: {mean_loss:.5}")
        tqdm.write(f"{mode}/Accuracy: {mean_acc:.5}")
        return mean_loss, mean_acc


def calculate_mean_std(train_dataset):
    channels_sum, channels_square_sum, batch = 0, 0, 0

    for data, label in train_dataset:
        # (rgb channel, w, h)
        channels_sum += torch.mean(data, dim=[1, 2])
        channels_square_sum += torch.mean(data ** 2, dim=[1, 2])
        batch += 1

    mean = channels_sum / batch
    std = (channels_square_sum / batch - mean * mean) ** 0.5
    return mean, std









